package uk.ac.rothamsted.utils;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;


import uk.ac.rothamsted.BlastData;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimaps;
import com.google.common.collect.SetMultimap;

public class OrthologClusterSet implements Serializable{
	/**
	 * 
	 */
	private static final long serialVersionUID = 4809124091588684395L;
	private final SetMultimap<String, OrthologCluster> membershipIndex = Multimaps.newSetMultimap(Maps.<String, Collection<OrthologCluster>> newHashMap(), DefaultSuppliers.set(OrthologCluster.class));
	private final Map<ArrayKey<String>, BlastData> allHits = new HashMap<ArrayKey<String>, BlastData> ();
	private final ListMultimap<String, BlastData> idToHitsIndex =  ArrayListMultimap.create();
	private final Map<ArrayKey<String>, OrthologCluster> seedOrthologs = new HashMap<ArrayKey<String>, OrthologCluster> ();
	private final Map<String, Integer> speciesSets = new HashMap<String, Integer>();

	//Comparator that can work on BlastData
	private final Comparator<BlastData> comp = new Comparator<BlastData>(){
		@Override
		public int compare(BlastData o1, BlastData o2) {
			if(o1.score == o2.score)
				return 0;
			else if(o1.score < o2.score)
				return -1;
			return 1;
		}
	};

	/**
	 * 
	 * @param allHits - all bidirectional hits in this blast set (AA, BB and AB) in one collection
	 * @param species - collections containing the complete set of id's belonging to particular species
	 */
	@SafeVarargs
	@SuppressWarnings({ "unchecked", "rawtypes" })
	public OrthologClusterSet(Collection<BlastData> allHits, Collection<String> ... species){
		for(BlastData bd : allHits){
			this.allHits.put(getSortedKey(bd.id1, bd.id2), bd);
			idToHitsIndex.put(bd.id1, bd);//index all hits to both sequence ids
			idToHitsIndex.put(bd.id2, bd);
		}
		Integer speciesSetId = 0;
		for(Collection<String> c :species){
			for(String id : c){
				speciesSets.put(id, speciesSetId);
			}
			speciesSetId++;
		}

		for(Entry<String, Collection<BlastData>> ent : idToHitsIndex.asMap().entrySet()){
			Collections.sort((List)ent.getValue(), comp);//sort the lists of all hits for particular sequence in ascending order
		}

	}

	/**
	 * Creates an empty cluster, containing just  the main pair of orthologs
	 * @param seedOrthologPair
	 * @return
	 */
	public OrthologCluster createEmptyCluster(BlastData seedOrthologPair){
		OrthologCluster cluster = new OrthologCluster(seedOrthologPair, this);
		seedOrthologs.put(getSortedKey(seedOrthologPair.id1, seedOrthologPair.id2), cluster);
		return cluster;
	}

	/**
	 * Creates a filled cluster, containing the main pair of orthologs and all valid inparalogs
	 * @param seedOrthologPair
	 * @return
	 */
	public OrthologCluster createFilledCluster(BlastData seedOrthologPair){
		OrthologCluster cluster = new OrthologCluster(seedOrthologPair, this);
		//Populate the cluster with paralogs

		List<BlastData> allHits = idToHitsIndex.get(seedOrthologPair.id1);// get list of all hits for first sequence
		if(allHits != null){
			for(int i = allHits.indexOf(seedOrthologPair); i < allHits.size(); i++){// as the list is sorted, all hits with the score higher than seedOrthologPair are above it in the list  
				if(isWithinSpeciesHit(allHits.get(i))){ // add to cluster if the hit is within the same species		
					cluster.add(allHits.get(i));
				}
			}
		}

		allHits = idToHitsIndex.get(seedOrthologPair.id2);
		if(allHits != null){
			for(int i = allHits.indexOf(seedOrthologPair); i < allHits.size(); i++){
				if(isWithinSpeciesHit(allHits.get(i))){

					cluster.add(allHits.get(i));
				}
			}
		}
		seedOrthologs.put(getSortedKey(seedOrthologPair.id1, seedOrthologPair.id2), cluster);
		return cluster;
	}

	void setMembership(String id, OrthologCluster cluster){
		membershipIndex.put(id, cluster);
	}

	void setMembership(Collection<String> ids, OrthologCluster cluster){
		for(String id : ids){
			membershipIndex.put(id, cluster);
		}
	}

	void removeMembership(String id, OrthologCluster cluster){
		membershipIndex.remove(id, cluster);
	}

	void removeMembership(Collection<String> ids, OrthologCluster cluster){
		for(String id : ids){
			membershipIndex.remove(id, cluster);	
		}
	}

	/**
	 * Returns clusters the supplied ids are assigned to			
	 * @param ids
	 * @return
	 */
	public Set<OrthologCluster> getClustersForIds(Collection<String> ids){	
		Set<OrthologCluster> result = new HashSet<OrthologCluster>();
		for(String id : ids){
			result.addAll(membershipIndex.get(id));
		}
		return result;
	}

	/**
	 * Returns true if the ids belong to the same species
	 * @param id1
	 * @param id2
	 * @return
	 */
	public boolean isSameSpecies(String id1, String id2){
		return speciesSets.get(id1).equals(speciesSets.get(id2));
	}

	public Integer getSpeciesId(String id){
		return speciesSets.get(id);
	}

	public void deleteCluster(OrthologCluster cluster){
		if(cluster.isInvalid()){
			return;
		}
		seedOrthologs.remove(getSortedKey(cluster.getMainPair().id1, cluster.getMainPair().id2));
		cluster.delete();
	}

	public OrthologCluster getCluster(String id1, String id2){
		return seedOrthologs.get(getSortedKey(id1, id2));
	}

	/**
	 * Returns a copy of all clusters currently held in this set (so be sure to dispose of this list when done!)
	 * @return
	 */
	public List<OrthologCluster> getClusterList(){
		List<OrthologCluster> clusterList = new ArrayList<OrthologCluster>(seedOrthologs.values());
		Collections.sort(clusterList);
		Collections.reverse(clusterList);
		return clusterList;
	}

	public ArrayKey<String> getSortedKey(String ... array){
		String []  ids =  Arrays.copyOf(array, array.length);
		Arrays.sort(ids);
		return new ArrayKey<String>(ids);
	}

	public boolean isWithinSpeciesHit(BlastData bd){
		return isSameSpecies(bd.id1, bd.id2);
	}

	public int size(){
		return seedOrthologs.size();
	}

	public Map<ArrayKey<String>, BlastData> getAllHits(){
		return allHits;
	}

	public List<Integer> clusterSize(){
		ArrayList<Integer> outSize = new ArrayList<Integer>();

		for(OrthologCluster eachCluster : seedOrthologs.values()){
			int capacity = outSize.size();
			while(capacity<=eachCluster.size()){
				for(int i = capacity; i< capacity+10; i++){
					outSize.add(i,0);
				}
				capacity = outSize.size();
			}
			int number = outSize.get(eachCluster.size()); 
			number++;
			outSize.set(eachCluster.size(), number);
		}
		return outSize;
	}

	public void printClusterSize(){
		List<Integer> sizes = this.clusterSize();
		int total = 0;
		for(int i=0; i<sizes.size(); i++){
			if(sizes.get(i)>0){
				System.out.println("Number of clusters with " + i + " paralogs is    " + sizes.get(i));
			}
			total +=i*sizes.get(i);
		}
		System.out.println("Total number of Paralogs found = " + total + "\n\n");
	}

	public boolean contains(OrthologCluster arg0){
		return seedOrthologs.containsKey(getSortedKey(arg0.getMainPair().id1, arg0.getMainPair().id2));
	}

	public String toString(){
		String output = "";
		List<OrthologCluster> cluster = this.getClusterList();
		for(int i = cluster.size()-1; i>=0;i--){
			if(cluster.get(i).isInvalid()){
				continue;
			}
			output += "#" + i + "  " + cluster.get(i) + "\n";
			i++;
		}
		return output;
	}

}

